#include "myTree.h"
#include <stack>
#include <queue>
#include <algorithm>

using namespace std;

// TreeNode类实现
TreeNode::TreeNode(char d, TreeNode* left, TreeNode* right, NodeTag lT, NodeTag rT)
    : data(d), leftChild(left), rightChild(right), lTag(lT), rTag(rT)
{
}

TreeNode::~TreeNode()
{
    // 析构函数，由MyTree负责管理内存
}

void TreeNode::printNode()
{
    cout << this->data;
}

// MyTree类实现
MyTree::MyTree()
    : root(nullptr), isThread(false)
{
}

// 根据先序序列构建二叉树的辅助函数
TreeNode* MyTree::buildTree(const char*& str)
{
    if (*str == '\0' || *str == '@')
    {
        if (*str == '@') str++;
        return nullptr;
    }
    
    TreeNode* node = new TreeNode(*str, nullptr, nullptr, Link, Link);
    str++;
    node->leftChild = buildTree(str);
    node->rightChild = buildTree(str);
    return node;
}

MyTree::MyTree(const char str[])
    : root(nullptr), isThread(false)
{
    const char* p = str;
    root = buildTree(p);
}

// 复制树的辅助函数
TreeNode* MyTree::copyTree(TreeNode* node)
{
    if (node == nullptr)
        return nullptr;
    
    TreeNode* newNode = new TreeNode(node->data, nullptr, nullptr, node->lTag, node->rTag);
    newNode->leftChild = copyTree(node->leftChild);
    newNode->rightChild = copyTree(node->rightChild);
    return newNode;
}

MyTree::MyTree(const MyTree& other)
    : root(nullptr), isThread(false)
{
    root = copyTree(other.root);
    isThread = false; // 复制后不是线索二叉树
}

// 销毁树的辅助函数
void MyTree::destroyTree(TreeNode* node)
{
    if (node == nullptr)
        return;
    
    // 如果是线索二叉树，需要先断开线索
    if (isThread)
    {
        if (node->lTag == Link && node->leftChild != nullptr)
            destroyTree(node->leftChild);
        if (node->rTag == Link && node->rightChild != nullptr)
            destroyTree(node->rightChild);
    }
    else
    {
        destroyTree(node->leftChild);
        destroyTree(node->rightChild);
    }
    
    delete node;
}

MyTree::~MyTree()
{
    destroyTree(root);
    root = nullptr;
}

// 先序遍历（非递归）
void MyTree::preOrderTraverse()
{
    if (root == nullptr)
        return;
    
    stack<TreeNode*> s;
    TreeNode* p = root;
    
    while (p != nullptr || !s.empty())
    {
        while (p != nullptr)
        {
            p->printNode();
            s.push(p);
            p = p->leftChild;
        }
        
        if (!s.empty())
        {
            p = s.top();
            s.pop();
            p = p->rightChild;
        }
    }
}

// 中序遍历（非递归，支持线索二叉树）
void MyTree::inOrderTraverse()
{
    if (root == nullptr)
        return;
    
    if (isThread)
    {
        // 线索二叉树的中序遍历
        TreeNode* p = root;
        // 找到最左节点
        while (p->leftChild != nullptr && p->lTag == Link)
            p = p->leftChild;
        
        while (p != nullptr)
        {
            p->printNode();
            
            // 如果有右线索，直接访问后继
            if (p->rTag == Thread)
            {
                p = p->rightChild;
            }
            else
            {
                // 否则进入右子树，找到最左节点
                p = p->rightChild;
                while (p != nullptr && p->lTag == Link)
                    p = p->leftChild;
            }
        }
    }
    else
    {
        // 普通二叉树的中序遍历
        stack<TreeNode*> s;
        TreeNode* p = root;
        
        while (p != nullptr || !s.empty())
        {
            while (p != nullptr)
            {
                s.push(p);
                p = p->leftChild;
            }
            
            if (!s.empty())
            {
                p = s.top();
                s.pop();
                p->printNode();
                p = p->rightChild;
            }
        }
    }
}

// 后序遍历（非递归）
void MyTree::postOrderTraverse()
{
    if (root == nullptr)
        return;
    
    stack<TreeNode*> s;
    TreeNode* p = root;
    TreeNode* lastVisited = nullptr;
    
    while (p != nullptr || !s.empty())
    {
        if (p != nullptr)
        {
            s.push(p);
            p = p->leftChild;
        }
        else
        {
            TreeNode* top = s.top();
            if (top->rightChild != nullptr && top->rightChild != lastVisited)
            {
                p = top->rightChild;
            }
            else
            {
                top->printNode();
                lastVisited = top;
                s.pop();
            }
        }
    }
}

// 计算叶子节点数
int MyTree::countLeaf()
{
    if (root == nullptr)
        return 0;
    
    int count = 0;
    stack<TreeNode*> s;
    s.push(root);
    
    while (!s.empty())
    {
        TreeNode* p = s.top();
        s.pop();
        
        if (p->leftChild == nullptr && p->rightChild == nullptr)
            count++;
        
        if (p->rightChild != nullptr)
            s.push(p->rightChild);
        if (p->leftChild != nullptr)
            s.push(p->leftChild);
    }
    
    return count;
}

// 计算树深度
int MyTree::countHeight()
{
    if (root == nullptr)
        return 0;
    
    queue<TreeNode*> q;
    q.push(root);
    int height = 0;
    
    while (!q.empty())
    {
        int size = q.size();
        height++;
        
        for (int i = 0; i < size; i++)
        {
            TreeNode* p = q.front();
            q.pop();
            
            if (p->leftChild != nullptr)
                q.push(p->leftChild);
            if (p->rightChild != nullptr)
                q.push(p->rightChild);
        }
    }
    
    return height;
}

// 判断是否为线索二叉树
bool MyTree::isThreadedTree()
{
    return isThread;
}

// 中序线索化的辅助函数
void MyTree::inOrderThreadingHelper(TreeNode* node, TreeNode*& pre)
{
    if (node == nullptr)
        return;
    
    inOrderThreadingHelper(node->leftChild, pre);
    
    if (node->leftChild == nullptr)
    {
        node->lTag = Thread;
        node->leftChild = pre;
    }
    else
    {
        node->lTag = Link;
    }
    
    if (pre != nullptr && pre->rightChild == nullptr)
    {
        pre->rTag = Thread;
        pre->rightChild = node;
    }
    else if (pre != nullptr)
    {
        pre->rTag = Link;
    }
    
    pre = node;
    
    inOrderThreadingHelper(node->rightChild, pre);
}

// 中序线索化
bool MyTree::inOrderThreading()
{
    if (root == nullptr || isThread)
        return false;
    
    TreeNode* pre = nullptr;
    inOrderThreadingHelper(root, pre);
    
    // 处理最后一个节点
    if (pre != nullptr && pre->rightChild == nullptr)
    {
        pre->rTag = Thread;
        pre->rightChild = nullptr; // 最后一个节点的右线索指向nullptr
    }
    
    isThread = true;
    return true;
}

// 定位节点
TreeNode& MyTree::locateNode(const char& v)
{
    static TreeNode nullNode('\0', nullptr, nullptr, Link, Link);
    
    if (root == nullptr)
        return nullNode;
    
    if (isThread)
    {
        // 线索二叉树的查找
        TreeNode* p = root;
        while (p->leftChild != nullptr && p->lTag == Link)
            p = p->leftChild;
        
        while (p != nullptr)
        {
            if (p->data == v)
                return *p;
            
            if (p->rTag == Thread)
            {
                p = p->rightChild;
            }
            else
            {
                p = p->rightChild;
                while (p != nullptr && p->lTag == Link)
                    p = p->leftChild;
            }
        }
    }
    else
    {
        // 普通二叉树的查找
        stack<TreeNode*> s;
        TreeNode* p = root;
        
        while (p != nullptr || !s.empty())
        {
            while (p != nullptr)
            {
                if (p->data == v)
                    return *p;
                s.push(p);
                p = p->leftChild;
            }
            
            if (!s.empty())
            {
                p = s.top();
                s.pop();
                if (p->data == v)
                    return *p;
                p = p->rightChild;
            }
        }
    }
    
    return nullNode;
}

// 获取前驱节点
TreeNode& MyTree::preNode(const TreeNode& node)
{
    static TreeNode nullNode('\0', nullptr, nullptr, Link, Link);
    
    if (!isThread)
        return nullNode;
    
    // 如果有左线索，直接返回
    if (node.lTag == Thread)
    {
        if (node.leftChild == nullptr)
            return nullNode;
        return *(node.leftChild);
    }
    
    // 否则，前驱是左子树的最右节点
    TreeNode* p = node.leftChild;
    if (p != nullptr)
    {
        while (p->rightChild != nullptr && p->rTag == Link)
            p = p->rightChild;
        return *p;
    }
    
    return nullNode;
}

// 获取后继节点
TreeNode& MyTree::nextNode(const TreeNode& node)
{
    static TreeNode nullNode('\0', nullptr, nullptr, Link, Link);
    
    if (!isThread)
        return nullNode;
    
    // 如果有右线索，直接返回
    if (node.rTag == Thread)
    {
        if (node.rightChild == nullptr)
            return nullNode;
        return *(node.rightChild);
    }
    
    // 否则，后继是右子树的最左节点
    TreeNode* p = node.rightChild;
    if (p != nullptr)
    {
        while (p->leftChild != nullptr && p->lTag == Link)
            p = p->leftChild;
        return *p;
    }
    
    return nullNode;
}

// HuffmanNode构造函数实现
HuffmanNode::HuffmanNode(int w, int v, bool leaf)
    : weight(w), value(v), left(nullptr), right(nullptr), isLeaf(leaf)
{
}

HuffmanTree::HuffmanTree(const int& n, const int arr[])
    : root(nullptr), nodeCount(n)
{
    // 创建节点数组，存储值和权重
    nodes = new pair<int, int>[n];
    for (int i = 0; i < n; i++)
    {
        nodes[i] = make_pair(arr[i], arr[i]);
    }
    
    // 构建霍夫曼树
    // 使用优先队列（最小堆），权重小的优先（值小的优先）
    struct Compare
    {
        bool operator()(HuffmanNode* a, HuffmanNode* b)
        {
            if (a->weight != b->weight)
                return a->weight > b->weight; // 权重小的优先
            return a->value > b->value; // 权重相同时，值小的优先
        }
    };
    
    priority_queue<HuffmanNode*, vector<HuffmanNode*>, Compare> pq;
    
    // 创建叶子节点，权重和值都等于arr[i]
    // 值越大，重要度越高，应该获得更短的编码
    for (int i = 0; i < n; i++)
    {
        HuffmanNode* node = new HuffmanNode(arr[i], arr[i], true);
        pq.push(node);
    }
    
    // 构建霍夫曼树
    while (pq.size() > 1)
    {
        HuffmanNode* first = pq.top();
        pq.pop();
        HuffmanNode* second = pq.top();
        pq.pop();
        
        // 创建新节点，权重为两个子节点之和
        int newWeight = first->weight + second->weight;
        int newValue = min(first->value, second->value);
        HuffmanNode* parent = new HuffmanNode(newWeight, newValue, false);
        
        // 偏小的值放入左子树，偏大的值放入右子树
        // first的权重更小（或相等时值更小），所以first放左子树
        parent->left = first;
        parent->right = second;
        
        pq.push(parent);
    }
    
    if (!pq.empty())
        root = pq.top();
}

HuffmanTree::~HuffmanTree()
{
    if (root != nullptr)
    {
        destroyHuffmanTree(root);
    }
    delete[] nodes;
}

void HuffmanTree::destroyHuffmanTree(HuffmanNode* node)
{
    if (node == nullptr)
        return;
    destroyHuffmanTree(node->left);
    destroyHuffmanTree(node->right);
    delete node;
}

// 生成编码的辅助函数
void HuffmanTree::generateCodes(HuffmanNode* node, string code, vector<pair<int, string>>& codes)
{
    if (node == nullptr)
        return;
    
    if (node->isLeaf)
    {
        codes.push_back(make_pair(node->value, code));
        return;
    }
    
    if (node->left != nullptr)
        generateCodes(node->left, code + "0", codes);
    if (node->right != nullptr)
        generateCodes(node->right, code + "1", codes);
}

void HuffmanTree::printHuffmanCodes()
{
    if (root == nullptr)
        return;
    
    vector<pair<int, string>> codes;
    generateCodes(root, "", codes);
    
    // 按节点值递减排序
    sort(codes.begin(), codes.end(), [](const pair<int, string>& a, const pair<int, string>& b) {
        return a.first > b.first;
    });
    
    for (const auto& p : codes)
    {
        cout << p.first << ":" << p.second << endl;
    }
}
